#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多目标优化引擎
实现NSGA-II和MOEA/D算法，提供帕累托前沿分析和多目标决策支持
"""

import numpy as np
import random
import math
from typing import List, Dict, Tuple, Any, Optional
from dataclasses import dataclass
import time
from copy import deepcopy

# 尝试相对导入，如果失败则使用绝对导入
try:
    from .models import (
        OptimizationProblem, ParetoSolution, MultiObjectiveResult, AlgorithmType,
        VariableType, OptimizationType
    )
except ImportError:
    from models import (
        OptimizationProblem, ParetoSolution, MultiObjectiveResult, AlgorithmType,
        VariableType, OptimizationType
    )

# ============================================================================
# NSGA-II算法实现
# ============================================================================

class NSGAII:
    """NSGA-II多目标优化算法"""
    
    def __init__(self, population_size: int = 100, max_generations: int = 250,
                 crossover_prob: float = 0.9, mutation_prob: float = 0.1):
        self.population_size = population_size
        self.max_generations = max_generations
        self.crossover_prob = crossover_prob
        self.mutation_prob = mutation_prob
        
    def optimize(self, problem: OptimizationProblem) -> MultiObjectiveResult:
        """执行NSGA-II优化"""
        start_time = time.time()
        
        # 初始化种群
        population = self._initialize_population(problem)
        
        # 评估初始种群
        for individual in population:
            individual.objective_values = self._evaluate_objectives(individual, problem)
        
        best_hypervolume = 0.0
        convergence_history = []
        
        # 主循环
        for generation in range(self.max_generations):
            # 生成子代
            offspring = self._generate_offspring(population, problem)
            
            # 合并父代和子代
            combined_population = population + offspring
            
            # 非支配排序
            fronts = self._fast_non_dominated_sort(combined_population)
            
            # 选择下一代
            population = self._environmental_selection(fronts, self.population_size)
            
            # 计算收敛指标
            pareto_front = fronts[0] if fronts else []
            hypervolume = self._calculate_hypervolume(pareto_front, problem)
            convergence_history.append(hypervolume)
            
            if hypervolume > best_hypervolume:
                best_hypervolume = hypervolume
        
        execution_time = time.time() - start_time
        
        # 构建结果
        final_pareto_front = fronts[0] if fronts else []
        spacing = self._calculate_spacing(final_pareto_front)
        
        return MultiObjectiveResult(
            problem_name=problem.name,
            algorithm_type=AlgorithmType.NSGA_II,
            pareto_front=final_pareto_front,
            execution_time=execution_time,
            iterations=self.max_generations,
            convergence_metrics={
                'final_hypervolume': best_hypervolume,
                'convergence_history': convergence_history
            },
            hypervolume=best_hypervolume,
            spacing=spacing,
            metadata={
                'population_size': self.population_size,
                'crossover_prob': self.crossover_prob,
                'mutation_prob': self.mutation_prob
            }
        )
    
    def _initialize_population(self, problem: OptimizationProblem) -> List[ParetoSolution]:
        """初始化种群"""
        population = []
        
        for _ in range(self.population_size):
            solution = {}
            for var in problem.variables:
                if var.variable_type == VariableType.CONTINUOUS:
                    value = random.uniform(var.lower_bound or 0, var.upper_bound or 1)
                elif var.variable_type == VariableType.INTEGER:
                    value = random.randint(int(var.lower_bound or 0), int(var.upper_bound or 10))
                elif var.variable_type == VariableType.BINARY:
                    value = random.choice([0, 1])
                else:
                    value = random.uniform(0, 1)
                solution[var.name] = value
            
            population.append(ParetoSolution(solution=solution, objective_values={}))
        
        return population
    
    def _evaluate_objectives(self, individual: ParetoSolution, problem: OptimizationProblem) -> Dict[str, float]:
        """评估目标函数"""
        objectives = {}
        
        for obj in problem.objectives:
            # 简化的目标函数评估（实际应用中需要解析表达式）
            if obj.expression:
                # 这里应该解析和计算表达式，暂时使用简化计算
                value = sum(individual.solution.values()) + random.uniform(-2, 2)
            else:
                # 默认计算方式
                value = sum(v * random.uniform(0.5, 1.5) for v in individual.solution.values())
            
            # 根据优化类型调整
            if obj.optimization_type == OptimizationType.MINIMIZE:
                value = -value
            
            objectives[obj.name] = value
        
        return objectives
    
    def _generate_offspring(self, population: List[ParetoSolution], problem: OptimizationProblem) -> List[ParetoSolution]:
        """生成子代"""
        offspring = []
        
        for _ in range(self.population_size):
            # 锦标赛选择
            parent1 = self._tournament_selection(population)
            parent2 = self._tournament_selection(population)
            
            # 交叉
            if random.random() < self.crossover_prob:
                child = self._crossover(parent1, parent2, problem)
            else:
                child = deepcopy(parent1)
            
            # 变异
            if random.random() < self.mutation_prob:
                child = self._mutation(child, problem)
            
            # 评估子代
            child.objective_values = self._evaluate_objectives(child, problem)
            offspring.append(child)
        
        return offspring
    
    def _tournament_selection(self, population: List[ParetoSolution], tournament_size: int = 2) -> ParetoSolution:
        """锦标赛选择"""
        tournament = random.sample(population, min(tournament_size, len(population)))
        return min(tournament, key=lambda x: (x.rank, -x.crowding_distance))
    
    def _crossover(self, parent1: ParetoSolution, parent2: ParetoSolution, problem: OptimizationProblem) -> ParetoSolution:
        """模拟二进制交叉（SBX）"""
        child_solution = {}
        eta_c = 20  # 交叉分布指数
        
        for var in problem.variables:
            var_name = var.name
            p1_val = parent1.solution[var_name]
            p2_val = parent2.solution[var_name]
            
            if random.random() <= 0.5:
                if abs(p1_val - p2_val) > 1e-14:
                    if p1_val < p2_val:
                        y1, y2 = p1_val, p2_val
                    else:
                        y1, y2 = p2_val, p1_val
                    
                    rand = random.random()
                    beta = 1.0 + (2.0 * (y1 - (var.lower_bound or 0)) / (y2 - y1))
                    alpha = 2.0 - beta ** (-(eta_c + 1.0))
                    
                    if rand <= (1.0 / alpha):
                        betaq = (rand * alpha) ** (1.0 / (eta_c + 1.0))
                    else:
                        betaq = (1.0 / (2.0 - rand * alpha)) ** (1.0 / (eta_c + 1.0))
                    
                    c1 = 0.5 * ((y1 + y2) - betaq * (y2 - y1))
                    c2 = 0.5 * ((y1 + y2) + betaq * (y2 - y1))
                    
                    # 边界检查
                    c1 = max(var.lower_bound or 0, min(var.upper_bound or 1, c1))
                    c2 = max(var.lower_bound or 0, min(var.upper_bound or 1, c2))
                    
                    child_solution[var_name] = c1 if random.random() <= 0.5 else c2
                else:
                    child_solution[var_name] = p1_val
            else:
                child_solution[var_name] = p1_val
        
        return ParetoSolution(solution=child_solution, objective_values={})
    
    def _mutation(self, individual: ParetoSolution, problem: OptimizationProblem) -> ParetoSolution:
        """多项式变异"""
        eta_m = 20  # 变异分布指数
        
        for var in problem.variables:
            var_name = var.name
            if random.random() < (1.0 / len(problem.variables)):
                y = individual.solution[var_name]
                yl = var.lower_bound or 0
                yu = var.upper_bound or 1
                
                delta1 = (y - yl) / (yu - yl)
                delta2 = (yu - y) / (yu - yl)
                
                rand = random.random()
                mut_pow = 1.0 / (eta_m + 1.0)
                
                if rand <= 0.5:
                    xy = 1.0 - delta1
                    val = 2.0 * rand + (1.0 - 2.0 * rand) * (xy ** (eta_m + 1.0))
                    deltaq = val ** mut_pow - 1.0
                else:
                    xy = 1.0 - delta2
                    val = 2.0 * (1.0 - rand) + 2.0 * (rand - 0.5) * (xy ** (eta_m + 1.0))
                    deltaq = 1.0 - val ** mut_pow
                
                y = y + deltaq * (yu - yl)
                y = max(yl, min(yu, y))
                individual.solution[var_name] = y
        
        return individual
    
    def _fast_non_dominated_sort(self, population: List[ParetoSolution]) -> List[List[ParetoSolution]]:
        """快速非支配排序"""
        fronts = [[]]
        
        for p in population:
            p.domination_count = 0
            p.dominated_solutions = []
            
            for q in population:
                if p.dominates(q):
                    p.dominated_solutions.append(q)
                elif q.dominates(p):
                    p.domination_count += 1
            
            if p.domination_count == 0:
                p.rank = 0
                fronts[0].append(p)
        
        i = 0
        while len(fronts[i]) > 0:
            next_front = []
            for p in fronts[i]:
                for q in p.dominated_solutions:
                    q.domination_count -= 1
                    if q.domination_count == 0:
                        q.rank = i + 1
                        next_front.append(q)
            i += 1
            fronts.append(next_front)
        
        return fronts[:-1]  # 移除最后一个空前沿
    
    def _calculate_crowding_distance(self, front: List[ParetoSolution]):
        """计算拥挤距离"""
        if len(front) <= 2:
            for solution in front:
                solution.crowding_distance = float('inf')
            return
        
        # 初始化拥挤距离
        for solution in front:
            solution.crowding_distance = 0
        
        # 对每个目标函数计算拥挤距离
        if front and front[0].objective_values:
            for obj_name in front[0].objective_values.keys():
                # 按目标函数值排序
                front.sort(key=lambda x: x.objective_values[obj_name])
                
                # 边界解设为无穷大
                front[0].crowding_distance = float('inf')
                front[-1].crowding_distance = float('inf')
                
                # 计算目标函数值范围
                obj_min = front[0].objective_values[obj_name]
                obj_max = front[-1].objective_values[obj_name]
                
                if obj_max - obj_min > 0:
                    for i in range(1, len(front) - 1):
                        distance = (front[i + 1].objective_values[obj_name] - 
                                  front[i - 1].objective_values[obj_name]) / (obj_max - obj_min)
                        front[i].crowding_distance += distance
    
    def _environmental_selection(self, fronts: List[List[ParetoSolution]], population_size: int) -> List[ParetoSolution]:
        """环境选择"""
        new_population = []
        
        for front in fronts:
            if len(new_population) + len(front) <= population_size:
                # 计算拥挤距离
                self._calculate_crowding_distance(front)
                new_population.extend(front)
            else:
                # 需要部分选择
                remaining = population_size - len(new_population)
                self._calculate_crowding_distance(front)
                front.sort(key=lambda x: -x.crowding_distance)
                new_population.extend(front[:remaining])
                break
        
        return new_population
    
    def _calculate_hypervolume(self, pareto_front: List[ParetoSolution], problem: OptimizationProblem) -> float:
        """计算超体积指标（简化版本）"""
        if not pareto_front:
            return 0.0
        
        # 简化的超体积计算
        reference_point = {}
        for obj in problem.objectives:
            # 使用最差值作为参考点
            worst_value = min(sol.objective_values[obj.name] for sol in pareto_front)
            reference_point[obj.name] = worst_value - 1.0
        
        total_volume = 0.0
        for solution in pareto_front:
            volume = 1.0
            for obj_name, obj_value in solution.objective_values.items():
                volume *= max(0, obj_value - reference_point[obj_name])
            total_volume += volume
        
        return total_volume
    
    def _calculate_spacing(self, pareto_front: List[ParetoSolution]) -> float:
        """计算间距指标"""
        if len(pareto_front) < 2:
            return 0.0
        
        distances = []
        for i, sol1 in enumerate(pareto_front):
            min_distance = float('inf')
            for j, sol2 in enumerate(pareto_front):
                if i != j:
                    distance = 0.0
                    for obj_name in sol1.objective_values.keys():
                        distance += (sol1.objective_values[obj_name] - sol2.objective_values[obj_name]) ** 2
                    distance = math.sqrt(distance)
                    min_distance = min(min_distance, distance)
            distances.append(min_distance)
        
        mean_distance = sum(distances) / len(distances)
        variance = sum((d - mean_distance) ** 2 for d in distances) / len(distances)
        return math.sqrt(variance)

# ============================================================================
# MOEA/D算法实现
# ============================================================================

class MOEAD:
    """MOEA/D多目标优化算法"""
    
    def __init__(self, population_size: int = 100, max_generations: int = 250,
                 neighborhood_size: int = 20, prob_neighbor_mating: float = 0.9):
        self.population_size = population_size
        self.max_generations = max_generations
        self.neighborhood_size = neighborhood_size
        self.prob_neighbor_mating = prob_neighbor_mating
        
    def optimize(self, problem: OptimizationProblem) -> MultiObjectiveResult:
        """执行MOEA/D优化"""
        start_time = time.time()
        
        # 生成权重向量
        weight_vectors = self._generate_weight_vectors(problem)
        
        # 初始化种群
        population = self._initialize_population(problem)
        
        # 计算邻域
        neighborhoods = self._calculate_neighborhoods(weight_vectors)
        
        # 初始化理想点
        ideal_point = self._initialize_ideal_point(population, problem)
        
        convergence_history = []
        
        # 主循环
        for generation in range(self.max_generations):
            for i in range(self.population_size):
                # 选择邻域或整个种群进行交配
                if random.random() < self.prob_neighbor_mating:
                    mating_pool = neighborhoods[i]
                else:
                    mating_pool = list(range(self.population_size))
                
                # 生成新解
                new_solution = self._generate_new_solution(population, mating_pool, problem)
                
                # 更新理想点
                self._update_ideal_point(new_solution, ideal_point)
                
                # 更新邻域解
                self._update_neighborhood(new_solution, i, neighborhoods[i], 
                                        population, weight_vectors, ideal_point)
            
            # 记录收敛信息
            hypervolume = self._calculate_hypervolume_moead(population, problem)
            convergence_history.append(hypervolume)
        
        execution_time = time.time() - start_time
        
        # 转换为帕累托解格式
        pareto_front = []
        for sol in population:
            pareto_solution = ParetoSolution(
                solution=sol.solution,
                objective_values=sol.objective_values
            )
            pareto_front.append(pareto_solution)
        
        # 执行非支配排序
        nsga2 = NSGAII()
        fronts = nsga2._fast_non_dominated_sort(pareto_front)
        final_pareto_front = fronts[0] if fronts else []
        
        spacing = nsga2._calculate_spacing(final_pareto_front)
        hypervolume = nsga2._calculate_hypervolume(final_pareto_front, problem)
        
        return MultiObjectiveResult(
            problem_name=problem.name,
            algorithm_type=AlgorithmType.MOEA_D,
            pareto_front=final_pareto_front,
            execution_time=execution_time,
            iterations=self.max_generations,
            convergence_metrics={
                'convergence_history': convergence_history,
                'final_hypervolume': hypervolume
            },
            hypervolume=hypervolume,
            spacing=spacing,
            metadata={
                'population_size': self.population_size,
                'neighborhood_size': self.neighborhood_size,
                'prob_neighbor_mating': self.prob_neighbor_mating
            }
        )
    
    def _generate_weight_vectors(self, problem: OptimizationProblem) -> List[List[float]]:
        """生成权重向量"""
        num_objectives = len(problem.objectives)
        weight_vectors = []
        
        if num_objectives == 2:
            # 二目标情况：均匀分布权重
            for i in range(self.population_size):
                w1 = i / (self.population_size - 1)
                w2 = 1 - w1
                weight_vectors.append([w1, w2])
        else:
            # 多目标情况：随机生成权重向量
            for _ in range(self.population_size):
                weights = [random.random() for _ in range(num_objectives)]
                total = sum(weights)
                weights = [w / total for w in weights]
                weight_vectors.append(weights)
        
        return weight_vectors
    
    def _initialize_population(self, problem: OptimizationProblem) -> List[ParetoSolution]:
        """初始化种群"""
        nsga2 = NSGAII()
        return nsga2._initialize_population(problem)
    
    def _calculate_neighborhoods(self, weight_vectors: List[List[float]]) -> List[List[int]]:
        """计算邻域"""
        neighborhoods = []
        
        for i, weight in enumerate(weight_vectors):
            distances = []
            for j, other_weight in enumerate(weight_vectors):
                if i != j:
                    distance = sum((w1 - w2) ** 2 for w1, w2 in zip(weight, other_weight))
                    distances.append((distance, j))
            
            distances.sort()
            neighborhood = [i] + [idx for _, idx in distances[:self.neighborhood_size - 1]]
            neighborhoods.append(neighborhood)
        
        return neighborhoods
    
    def _initialize_ideal_point(self, population: List[ParetoSolution], problem: OptimizationProblem) -> Dict[str, float]:
        """初始化理想点"""
        ideal_point = {}
        
        # 评估初始种群
        nsga2 = NSGAII()
        for individual in population:
            individual.objective_values = nsga2._evaluate_objectives(individual, problem)
        
        # 找到每个目标的最优值
        for obj in problem.objectives:
            best_value = max(sol.objective_values[obj.name] for sol in population)
            ideal_point[obj.name] = best_value
        
        return ideal_point
    
    def _generate_new_solution(self, population: List[ParetoSolution], mating_pool: List[int], 
                             problem: OptimizationProblem) -> ParetoSolution:
        """生成新解"""
        # 随机选择两个父代
        parent_indices = random.sample(mating_pool, min(2, len(mating_pool)))
        parent1 = population[parent_indices[0]]
        parent2 = population[parent_indices[1]] if len(parent_indices) > 1 else parent1
        
        # 使用NSGA-II的交叉和变异操作
        nsga2 = NSGAII()
        child = nsga2._crossover(parent1, parent2, problem)
        child = nsga2._mutation(child, problem)
        
        # 评估新解
        child.objective_values = nsga2._evaluate_objectives(child, problem)
        
        return child
    
    def _update_ideal_point(self, solution: ParetoSolution, ideal_point: Dict[str, float]):
        """更新理想点"""
        for obj_name, obj_value in solution.objective_values.items():
            if obj_value > ideal_point[obj_name]:
                ideal_point[obj_name] = obj_value
    
    def _update_neighborhood(self, new_solution: ParetoSolution, index: int, neighborhood: List[int],
                           population: List[ParetoSolution], weight_vectors: List[List[float]],
                           ideal_point: Dict[str, float]):
        """更新邻域解"""
        for neighbor_idx in neighborhood:
            # 计算切比雪夫距离
            current_fitness = self._tchebycheff_fitness(population[neighbor_idx], 
                                                      weight_vectors[neighbor_idx], ideal_point)
            new_fitness = self._tchebycheff_fitness(new_solution, 
                                                   weight_vectors[neighbor_idx], ideal_point)
            
            # 如果新解更好，则替换
            if new_fitness < current_fitness:
                population[neighbor_idx] = deepcopy(new_solution)
    
    def _tchebycheff_fitness(self, solution: ParetoSolution, weight_vector: List[float], 
                           ideal_point: Dict[str, float]) -> float:
        """计算切比雪夫适应度"""
        max_value = 0.0
        
        obj_names = list(solution.objective_values.keys())
        for i, obj_name in enumerate(obj_names):
            if i < len(weight_vector):
                value = weight_vector[i] * abs(solution.objective_values[obj_name] - ideal_point[obj_name])
                max_value = max(max_value, value)
        
        return max_value
    
    def _calculate_hypervolume_moead(self, population: List[ParetoSolution], problem: OptimizationProblem) -> float:
        """计算MOEA/D的超体积"""
        nsga2 = NSGAII()
        return nsga2._calculate_hypervolume(population, problem)

# ============================================================================
# 多目标优化引擎主类
# ============================================================================

class MultiObjectiveEngine:
    """多目标优化引擎"""
    
    def __init__(self):
        self.algorithms = {
            AlgorithmType.NSGA_II: NSGAII,
            AlgorithmType.MOEA_D: MOEAD
        }
    
    def optimize(self, problem: OptimizationProblem, algorithm_type: AlgorithmType = AlgorithmType.NSGA_II,
                 **kwargs) -> MultiObjectiveResult:
        """执行多目标优化"""
        # 检查是否为多目标问题
        if len(problem.objectives) < 2:
            raise ValueError("多目标优化至少需要2个目标函数")
        
        if algorithm_type not in self.algorithms:
            raise ValueError(f"不支持的算法类型: {algorithm_type}")
        
        algorithm_class = self.algorithms[algorithm_type]
        
        # 根据算法类型过滤参数
        if algorithm_type == AlgorithmType.NSGA_II:
            # NSGA-II支持的参数
            valid_params = {k: v for k, v in kwargs.items() 
                          if k in ['population_size', 'max_generations', 'crossover_prob', 'mutation_prob']}
        elif algorithm_type == AlgorithmType.MOEA_D:
            # MOEA/D支持的参数
            valid_params = {k: v for k, v in kwargs.items() 
                          if k in ['population_size', 'max_generations', 'neighborhood_size', 'prob_neighbor_mating']}
        else:
            valid_params = kwargs
        
        algorithm = algorithm_class(**valid_params)
        
        return algorithm.optimize(problem)
    
    def analyze_pareto_front(self, result: MultiObjectiveResult) -> Dict[str, Any]:
        """分析帕累托前沿"""
        if not result.pareto_front:
            return {"error": "帕累托前沿为空"}
        
        analysis = {
            "solution_count": len(result.pareto_front),
            "hypervolume": result.hypervolume,
            "spacing": result.spacing,
            "objective_ranges": {},
            "best_solutions": []
        }
        
        # 分析目标函数范围
        if result.pareto_front[0].objective_values:
            for obj_name in result.pareto_front[0].objective_values.keys():
                values = [sol.objective_values[obj_name] for sol in result.pareto_front]
                analysis["objective_ranges"][obj_name] = {
                    "min": min(values),
                    "max": max(values),
                    "mean": sum(values) / len(values)
                }
        
        # 获取最佳解
        analysis["best_solutions"] = result.get_best_solutions(3)
        
        return analysis
    
    def generate_decision_support(self, result: MultiObjectiveResult, preferences: Dict[str, float] = None) -> Dict[str, Any]:
        """生成决策支持信息"""
        if not result.pareto_front:
            return {"error": "无可用解"}
        
        support = {
            "recommended_solutions": [],
            "trade_off_analysis": {},
            "sensitivity_analysis": {}
        }
        
        # 基于偏好推荐解
        if preferences:
            scored_solutions = []
            for sol in result.pareto_front:
                score = 0.0
                for obj_name, weight in preferences.items():
                    if obj_name in sol.objective_values:
                        score += weight * sol.objective_values[obj_name]
                scored_solutions.append((score, sol))
            
            scored_solutions.sort(key=lambda x: x[0], reverse=True)
            support["recommended_solutions"] = [sol for _, sol in scored_solutions[:5]]
        else:
            support["recommended_solutions"] = result.get_best_solutions(5)
        
        # 权衡分析
        if len(result.pareto_front) > 1:
            obj_names = list(result.pareto_front[0].objective_values.keys())
            if len(obj_names) >= 2:
                correlations = {}
                for i, obj1 in enumerate(obj_names):
                    for j, obj2 in enumerate(obj_names[i+1:], i+1):
                        values1 = [sol.objective_values[obj1] for sol in result.pareto_front]
                        values2 = [sol.objective_values[obj2] for sol in result.pareto_front]
                        
                        # 简单相关性计算
                        mean1, mean2 = sum(values1)/len(values1), sum(values2)/len(values2)
                        numerator = sum((v1-mean1)*(v2-mean2) for v1,v2 in zip(values1,values2))
                        denom1 = sum((v1-mean1)**2 for v1 in values1)**0.5
                        denom2 = sum((v2-mean2)**2 for v2 in values2)**0.5
                        
                        if denom1 > 0 and denom2 > 0:
                            correlation = numerator / (denom1 * denom2)
                            correlations[f"{obj1}_vs_{obj2}"] = correlation
                
                support["trade_off_analysis"] = correlations
        
        return support